import os, sys
from typing import Dict

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 4)))
sys.path.insert(0, parent_path)

from base64 import b64encode
import traceback
import numpy as np
import socket, json, cv2, yaml, time
from typing import *
from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.PIL_draw import pic_text
from core.utils.visualize import get_color_map_list

class RTMPProcessor(BaseProcessor):
    """直播流检测后处理方案
    """
    def __init__(self, 
                 keys, 
                 Custom_cls_map,
                 Custom_target_cls,
                 Custom_stride_time,
                 Custom_continuous_threshold
                 ) -> None:
        self.keys = keys
        if isinstance(Custom_target_cls,str):
            with open(Custom_target_cls,"r") as f:
                self.target_cls = set(yaml.safe_load(f)["label_list"])
        else:
            self.target_cls=set(Custom_target_cls)
        if isinstance(Custom_cls_map,str):
            with open(Custom_cls_map,"r") as f:
                self.cls_map = {x:y for x,y in enumerate(yaml.safe_load(f)["label_list"])}
        else:
            self.cls_map = Custom_cls_map
        self.color_list = get_color_map_list(80)
        self.stride_time = Custom_stride_time
        self.last_up_time = time.time()
        self.continuous_threshold = Custom_continuous_threshold
        self.continuous_count = 0

    def __call__(self, data:Dict) -> Dict:
        
        img = data[self.keys["in"]][0].copy()
        push_path = data["SteamPush_url"] if data.get("SteamPush_url",False) else ""
        have_target = False
        draw_thickness = min(img.shape[:2]) // 320
        det_result = data[self.keys["det"]]
        objects = {}
        for idx_, bboxes in enumerate(det_result["boxes"][0]):
            cls_id, score, x1,y1,x2,y2 = bboxes
            cls_id = int(cls_id)
            if cls_id not in self.cls_map:
                continue
            if self.cls_map[cls_id] in self.target_cls:
                have_target = True
            object_ = {
                "cls_name": self.cls_map[cls_id],
                "box": [cls_id, score, x1, y1, x2, y2],
            }
            objects[idx_] = object_
            img = cv2.rectangle(img,(int(x1), int(y1)),(int(x2),int(y2)),color=self.color_list[cls_id],thickness=draw_thickness)
            text = self.cls_map[cls_id]
            cv2.putText(img,text,(int(x1), int(y1)),cv2.FONT_HERSHEY_COMPLEX_SMALL,1,self.color_list[cls_id])
            # img = pic_text(img,text,(int(x1), int(y1)),self.color_list[cls_id],draw_thickness*10)
        data[self.keys["out_vis"]] = img
        # 数据上报
        if have_target & (time.time() - self.last_up_time > self.stride_time):
            self.continuous_count+=1
            if self.continuous_count>=self.continuous_threshold:
                self.last_up_time = time.time()
                img_byte = cv2.imencode(".jpg",img)[1].tobytes()
                img_string = b64encode(img_byte).decode("utf-8")     # 图片64位编码
                data[self.keys["out_obj"]] = {"img":img_string,"object":objects,"push_path":push_path}
                self.continuous_count = 0
        else:
            self.continuous_count = 0
            
        return data
    
    def close(self):
        self.status = 0

class FTPProcessor(BaseProcessor):
    """FTP检测后处理方案
    """
    def __init__(self, 
                 keys,
                 Custom_cls_map,
                 Custom_target_cls,
                 ) -> None:
        self.keys = keys
        if isinstance(Custom_target_cls,str):
            with open(Custom_target_cls,"r") as f:
                self.target_cls = set(yaml.safe_load(f)["label_list"])
        else:
            self.target_cls=set(Custom_target_cls)
        if isinstance(Custom_cls_map,str):
            with open(Custom_cls_map,"r") as f:
                self.cls_map = {x:y for x,y in enumerate(yaml.safe_load(f)["label_list"])}
        else:
            self.cls_map = Custom_cls_map
        self.color_list = get_color_map_list(80)
        
    def init_check(self):
        assert "det" in self.keys
        super().init_check()
    
    def __call__(self, data:Dict) -> Dict:
        
        batch_img = data[self.keys["in"]].copy()
        det_result = data[self.keys["det"]]
        batch_objects = []
        batch_vis = []
        # start_idx = 0
        for img_idx, bboxes in enumerate(det_result["boxes"]):
            draw_thickness = min(batch_img[img_idx].shape[:2]) // 320
            img_objects = []
            img = batch_img[img_idx].copy()
            for bbox in bboxes:
                cls_id, score, x1,y1,x2,y2 = bbox
                cls_id = int(cls_id)
                object = {
                    "box": [cls_id, score, x1, y1, x2, y2],
                    "cls_name": self.cls_map[cls_id]
                }
                img_objects.append(object)
                # vis
                img = cv2.rectangle(img,(int(x1), int(y1)),(int(x2),int(y2)),color=self.color_list[cls_id],thickness=draw_thickness)
                text = self.cls_map[cls_id]
                img = pic_text(img,text,(int(x1), int(y1)),self.color_list[cls_id],draw_thickness*10)
            batch_objects.append(img_objects)
            batch_vis.append(img)
            # result["vis"].append(img)
            # result["objects"][data["data"][img_idx]] = {
            #     "download":True,
            #     "img_size":[batch_img[img_idx].shape[1],batch_img[img_idx].shape[0]],
            #     "object":img_objects
            # }
        
        data[self.keys["out_vis"]] = batch_vis
        data[self.keys["out_obj"]] = batch_objects
            
        return data

        
class WaterPollutionProcessor(BaseProcessor):
    """FTP检测后处理方案
    """
    def __init__(self, 
                 keys,
                 Custom_cls_map,
                 Custom_target_cls,
                 ) -> None:
        self.keys = keys
        if isinstance(Custom_target_cls,str):
            with open(Custom_target_cls,"r") as f:
                self.target_cls = set(yaml.safe_load(f)["label_list"])
        else:
            self.target_cls=set(Custom_target_cls)
        if isinstance(Custom_cls_map,str):
            with open(Custom_cls_map,"r") as f:
                self.cls_map = {x:y for x,y in enumerate(yaml.safe_load(f)["label_list"])}
        else:
            self.cls_map = Custom_cls_map
        self.color_list = get_color_map_list(80)
        
    def init_check(self):
        assert "det" in self.keys
        super().init_check()
    
    def __call__(self, data:Dict) -> Dict:
        
        batch_img = data[self.keys["in"]].copy()
        det_result = data[self.keys["det"]]
        result = {"vis":[],"objects":[]}
        # start_idx = 0
        for idx, bboxes in enumerate(det_result["boxes"]):
            draw_thickness = min(batch_img[idx].shape[:2]) // 320
            # bboxes = det_result["boxes"][start_idx:start_idx+box_num]
            # start_idx+=box_num
            objects = {}
            for idx_,boxes in enumerate(bboxes):
                cls_id, score, x1,y1,x2,y2 = boxes
                cls_id = int(cls_id)
                object_ = {
                    "box": [x1,y1,x2,y2],
                    "cls_name": self.cls_map[int(cls_id)]
                }
                objects[idx_] = object_
                img = cv2.rectangle(batch_img[idx],(int(x1), int(y1)),(int(x2),int(y2)),color=self.color_list[cls_id],thickness=draw_thickness)
                text = self.cls_map[cls_id]
                img = pic_text(img,text,(int(x1), int(y1)),self.color_list[cls_id],draw_thickness*10)
            result["vis"].append(img)
            result["objects"].append(objects)
        
        data[self.keys["out_vis"]] = result["vis"]
        data[self.keys["out_obj"]] = result["objects"]
            
        return data

    def close(self):
        self.status = 0